package com.wfm.leetcode.editor.cn;

import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;

/**
 * 查找和最小的 K 对数字
 * 2025-03-03 14:55:36
 * 已经选的数对假设为(ai,bi),那么下一个待选的是(ai+1,bi)(ai,bi+1)。
 * 如果每次都将已选的数对加到待选堆中，可能会出现重复，
 * 因此可以将nums1的前k个索引数对加入到队列中，每次从队列中取出元素(x,y),然后只需要将nums2的索引增加即可。
 * 时间复杂度：O(klogk)，k是数对的数目
 * 空间复杂度：O(k)，
 */
class FindKPairsWithSmallestSums {

    //leetcode submit region begin(Prohibit modification and deletion)
class Solution {
    public List<List<Integer>> kSmallestPairs(int[] nums1, int[] nums2, int k) {
        PriorityQueue<int[]> pq = new PriorityQueue<>(k, (o1, o2)->{
            return nums1[o1[0]] + nums2[o1[1]] - nums1[o2[0]] - nums2[o2[1]];
        });
        List<List<Integer>> ans = new ArrayList<>();
        int m = nums1.length;
        int n = nums2.length;
        for (int i = 0; i < Math.min(m, k); i++) {
            pq.offer(new int[]{i,0});
        }
        while (k-- > 0 && !pq.isEmpty()) {
            int[] idxPair = pq.poll();
            List<Integer> list = new ArrayList<>();
            list.add(nums1[idxPair[0]]);
            list.add(nums2[idxPair[1]]);
            ans.add(list);
            if (idxPair[1] + 1 < n) {
                pq.offer(new int[]{idxPair[0], idxPair[1] + 1});
            }
        }

        return ans;
    }
}
//leetcode submit region end(Prohibit modification and deletion)

    public static void main(String[] args) {
        Solution solution = new FindKPairsWithSmallestSums().new Solution();
        
    }
}